from abc import ABC, abstractmethod
from typing import Any, Dict, List, Optional
from pydantic import BaseModel


class Document(BaseModel):
    """文档基类"""
    content: str
    metadata: Dict[str, Any]
    score: Optional[float] = None


class SearchResult(BaseModel):
    """搜索结果基类"""
    documents: List[Document]
    source: str
    metadata: Dict[str, Any]


class BaseModule(ABC):
    """模块基类"""
    @abstractmethod
    def initialize(self, config: Dict[str, Any]) -> None:
        """初始化模块"""
        pass

    @abstractmethod
    def validate_config(self, config: Dict[str, Any]) -> bool:
        """验证配置"""
        pass


class BaseSearchStrategy(BaseModule):
    """搜索策略基类"""
    @abstractmethod
    async def search(self, query: str, **kwargs) -> SearchResult:
        """执行搜索"""
        pass


class BasePredictionModule(BaseModule):
    """预测模块基类"""
    @abstractmethod
    async def analyze_intent(self, query: str) -> Dict[str, Any]:
        """分析查询意图"""
        pass

    @abstractmethod
    async def generate_context(self, query: str, intent: Dict[str, Any]) -> str:
        """生成上下文"""
        pass

    @abstractmethod
    async def rewrite_query(self, query: str, context: str) -> List[str]:
        """重写查询"""
        pass


class BaseMemoryModule(BaseModule):
    """记忆模块基类"""
    @abstractmethod
    async def store(self, session_id: str, data: Dict[str, Any]) -> None:
        """存储对话历史"""
        pass

    @abstractmethod
    async def retrieve(self, session_id: str, k: int = 5) -> List[Dict[str, Any]]:
        """检索对话历史"""
        pass

    @abstractmethod
    async def update(self, session_id: str, data: Dict[str, Any]) -> None:
        """更新对话历史"""
        pass


class BaseFusionModule(BaseModule):
    """融合模块基类"""
    @abstractmethod
    async def expand_query(self, query: str) -> List[str]:
        """扩展查询"""
        pass

    @abstractmethod
    async def merge_results(self, results: List[SearchResult]) -> SearchResult:
        """合并结果"""
        pass

    @abstractmethod
    async def rerank(self, query: str, results: SearchResult) -> SearchResult:
        """重排序结果"""
        pass


class BaseRouterModule(BaseModule):
    """路由模块基类"""
    @abstractmethod
    async def route(self, query: str, context: Dict[str, Any]) -> List[str]:
        """路由查询"""
        pass

    @abstractmethod
    async def aggregate_results(self, results: List[SearchResult]) -> SearchResult:
        """聚合结果"""
        pass


class BaseTaskAdapter(BaseModule):
    """任务适配器基类"""
    @abstractmethod
    async def adapt(self, query: str, context: Dict[str, Any]) -> Dict[str, Any]:
        """适配任务"""
        pass

    @abstractmethod
    async def format_output(self, result: Any) -> Dict[str, Any]:
        """格式化输出"""
        pass


class ModuleRegistry:
    """模块注册器"""
    _modules: Dict[str, Dict[str, Any]] = {
        "search": {},
        "prediction": {},
        "memory": {},
        "fusion": {},
        "router": {},
        "adapter": {}
    }

    @classmethod
    def register(cls, module_type: str, name: str, module_class: Any) -> None:
        """注册模块"""
        if module_type not in cls._modules:
            raise ValueError(f"Unknown module type: {module_type}")
        cls._modules[module_type][name] = module_class

    @classmethod
    def get(cls, module_type: str, name: str) -> Any:
        """获取模块"""
        if module_type not in cls._modules:
            raise ValueError(f"Unknown module type: {module_type}")
        if name not in cls._modules[module_type]:
            raise ValueError(f"Unknown module name: {name}")
        return cls._modules[module_type][name]

    @classmethod
    def list_modules(cls, module_type: str) -> List[str]:
        """列出模块"""
        if module_type not in cls._modules:
            raise ValueError(f"Unknown module type: {module_type}")
        return list(cls._modules[module_type].keys()) 